import torch
from torch.nn.modules.module import Module


'''
  编码层
  输入 x1（n,n,in_feature,modes）   x2(n,n,out_feature,modes)
  输出（n，n，1）
  示例：
    x1 = torch.rand([100,100,4,4])
    x2 = torch.rand([100,100,5,6])
    model = Encode()
    y = model(x1,x2)
    print(y.shape)
'''

class Encode(Module):  
    def __init__(self):
        super(Encode,self).__init__()
    def forward(self,x1,x2):   #x1(batch,n,n,in_feature,modes)   x2(batch,n,n,out_feature,modes)    out(batch,n,n,1)
        x1_merged_tensor = x1.view(x1.shape[0], x1.shape[1], x1.shape[2] , x1.shape[3] * x1.shape[4])
        x2_merged_tensor = x2.view(x2.shape[0], x2.shape[1], x2.shape[2] , x2.shape[3] * x2.shape[4])
        combined_tensor = torch.cat((x1_merged_tensor, x2_merged_tensor), dim=3)
        output = combined_tensor.sum(dim=3, keepdim=True)
        return output 